from clickhouse_driver import Client
from common.variables import TimeFrameType
from sqlalchemy import create_engine, Engine
import toml


config_file: str = "config.toml"
config = toml.load(config_file)
is_prod = config["is_prod"]


def get_db_engine(db: str) -> Engine:
    engine = create_engine(
        f"clickhouse://{config['clickhouse']['user']}:{config['clickhouse']['pw']}@localhost:8123/{db}"
    )
    return engine


def get_db_driver_client(db: str) -> Client:
    # 连接 ClickHouse 服务器
    client = Client(
        host="localhost",
        port=9000,
        user=config["clickhouse"]["user"],
        password=config["clickhouse"]["pw"],
        database=db,
    )
    return client


def get_kafka_db() -> str:
    return "kafka_coins" if is_prod else "test_kafka_coins"


def get_db() -> str:
    return "coins" if is_prod else "test_coins"


def get_kline_tb(timeframe: TimeFrameType) -> str:
    return f"klines_{timeframe}"


def get_factor_tb(timeframe: TimeFrameType) -> str:
    return f"factor_{timeframe}"


def get_quant_db():
    return "quant"


def get_feature_tb(timeframe: TimeFrameType) -> str:
    return f"feature_{timeframe}"
